import math
import skimage
import numpy as np
import os
import matplotlib.pyplot as plt

import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = '0'  # 定位cuda设备为GPU0（需要根据情况修稿）
config = tf.compat.v1.ConfigProto()
# config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9  # 最大可申请显存比例
config.gpu_options.allow_growth = True  # 允许动态分配显存
tf.config.experimental.set_memory_growth = True
sess = tf.compat.v1.Session(config=config)

# from keras import backend
from keras.layers import *
from keras.models import load_model
# from skimage.measure import compare_psnr
from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from unet3 import cross_entropy_balanced
from mpl_toolkits.mplot3d import Axes3D
import mpl_toolkits.axisartist as ast
import os
from faultconnect import datainput
import cigvis
from cigvis import colormap

from medpy import metric
from keras import backend as K
from display3D import displayfault3d, displaytrain3d, display_seis_3d, display3d, displayfaultF33d
# import moxing as mox
import mayavi.mlab as mlab
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import PrecisionRecallDisplay
from utils import DataGenerator

pngDir = './png/'

model = load_model('GIE_Net/U_Net_1_50-30.hdf5',
                   # 'impd.hdf5',
                   custom_objects={
                      'cross_entropy_balanced': cross_entropy_balanced
                   }

                   )

model1 = load_model('GIE_Net/GIE_Net_1_50-36.hdf5',
                    # 'impd.hdf5',
                    custom_objects={
                       'cross_entropy_balanced': cross_entropy_balanced
                    }

                    )


def main():
    # goTrainTest(60)

    # goValidTest(model1)

   
    # 运行实际工区
    trainF3()
 


def goTrainTest(idnum):
    seismPath = "./data/train/seis/"
    faultPath = "./data/train/fault/"
    n1, n2, n3 = 128, 128, 128
    dk = idnum
    gx = np.load(seismPath + str(dk) + '.npy')
    fx = np.load(faultPath + str(dk) + '.npy')
    # gx = np.reshape(gx, (n1, n2, n3))
    # fx = np.reshape(fx, (n1, n2, n3))
    # display3d(gx)
    # display3d(fx)
    gm = np.mean(gx)
    gs = np.std(gx)
    gx = gx - gm
    gx = gx / gs
    gx = np.transpose(gx)
    fx = np.transpose(fx)

    displaytrain3d(gx=gx, fx=fx)
    fp = model.predict(np.reshape(gx, (1, n1, n2, n3, 1)), verbose=1)
    fp = fp[0, :, :, :, 0]
    # 将数据降维，要不然无法绘制其三维图
    # 获取图像的二维坐标，下面是一个三维图像，分别取其x,y,z面的数据，画出其平面图。
    gx1 = gx[61, :, :]
    fx1 = fx[61, :, :]
    fp1 = fp[61, :, :]
    gx2 = gx[:, 21, :]
    fx2 = fx[:, 21, :]
    fp2 = fp[:, 21, :]
    gx3 = gx[:, :, 21]
    fx3 = fx[:, :, 21]
    fp3 = fp[:, :, 21]


def goValidTest(model):
    # 路径
    seismPath = "./data/validation/seis/"
    faultPath = "./data/validation/fault/"
    n1, n2, n3 = 128, 128, 128
    dk = 9            # 3
    gx = np.fromfile(seismPath + str(dk) + '.dat', dtype=np.single)
    fx = np.fromfile(faultPath + str(dk) + '.dat', dtype=np.single)
    gx = np.reshape(gx, (n1, n2, n3))
    fx = np.reshape(fx, (n1, n2, n3))

    gm = np.mean(gx)
    gs = np.std(gx)
    gx = gx - gm
    gx = gx / gs
    '''
    gmin = np.min(gx)
    gmax = np.max(gx)
    gx = gx-gmin
    gx = gx/(gmax-gmin)
    '''
    gx = np.transpose(gx)
    fx = np.transpose(fx)

    # displayfault3d(gx=gx, fx=fx)
    fp = model.predict(np.reshape(gx, (1, n1, n2, n3, 1)), verbose=1)
    fp = fp[0, :, :, :, 0]
    # 将数据降维，要不然无法绘制其三维图
    # 获取图像的二维坐标，下面是一个三维图像，分别取其x,y,z面的数据，画出其平面图。
    gx1 = gx[61, :, :]
    fx1 = fx[61, :, :]
    fp1 = fp[61, :, :]
    gx2 = gx[:, 21, :]
    fx2 = fx[:, 21, :]
    fp2 = fp[:, 21, :]
    gx3 = gx[:, :, 21]
    fx3 = fx[:, :, 21]
    fp3 = fp[:, :, 21]
   plot2d(gx1, fx1, fp1, at=1, png='f3d/fp1')

    # fp1 = np.rint(fp)
    # fx1 = fx

    # fp1 = np.array(fp1)
    # fx1 = np.array(fx1)
    # res = calculate_metrics(fp1, fx1, "iou")  # iou， dice_coefficient， accuracy， precision， recall，
    # # sensitivity， f1, specificity
    # print('iou:{:.4f}'.format(res))
    # print('dice_coefficient:{:.4f}'.format(calculate_metrics(fp1, fx1, "dice_coefficient")))
    # print('accuracy:{:.4f}'.format(calculate_metrics(fp1, fx1, "accuracy")))
    # print('precision:{:.4f}'.format(calculate_metrics(fp1, fx1, "precision")))
    # print('recall:{:.4f}'.format(calculate_metrics(fp1, fx1, "recall")))
    # print('specificity:{:.4f}'.format(calculate_metrics(fp1, fx1, "specificity")))
    # print('f1:{:.4f}'.format(calculate_metrics(fp1, fx1, "f1")))
   
    # rec = metric.binary.recall(fp1, fx1)
    # print('recall:{}'.format(rec))
    # pre = metric.binary.dc(fp1, fx1)
    # print('dc:{}'.format(pre))


def trainF3():
    seismPath = "./data/prediction/f3d/"
    # kerry3d_npy = './data/kerry/Kerry3D.npz'
    n3, n2, n1 = 512, 384, 128

    gx = np.fromfile(seismPath + 'gxl.dat', dtype=np.single)
    fx = np.fromfile(seismPath + 'fpx.dat', dtype=np.single)

    gx = np.reshape(gx, (n3, n2, n1))
    fx = np.reshape(fx, (n3, n2, n1))
    gm = np.mean(gx)
    gs = np.std(gx)
    gx = gx - gm
    gx = gx / gs     # 数据标准化，保证网络更好的收敛，提高网络模型的精度
    '''
    gmin = np.min(gx)
    gmax = np.max(gx)
    gx = gx-gmin
    gx = gx/(gmax-gmin)
    '''
    gx = np.transpose(gx)
    fx = np.transpose(fx)


    with tf.device("/cpu"):
        # displayfault3d(gx=gx, fx=fx)
        print('ss')
        fp = model1.predict(np.reshape(gx, (1, 128, 384, 512, 1)),  verbose=1)

        fp = fp[0, :, :, :, 0]
        # 三维显示
        displayfaultF33d(gx=gx, fx=fp)
        # # 格式化操作
        # # fp = np.where(fp<0.2, 0, fp)
        # # fp = np.rint(fp)
        # gx = np.transpose(gx)
        # fx = np.transpose(fx)
        # fp = np.transpose(fp)
        fl1 = gx[90, :, :]
        fx1 = fx[90, :, :]
        fp1 = fp[90, :, :]
        separate_plot2d(fl1, fx1, fp1, at=1, png='f3d/fp1')
        # plot2d(fl1, fx1, fp1, at=1, png='f3d/fp1')

def calculate_metrics(predict_image, gt_image, evaluate):
    # 将图像转换为二进制数组
    # predict_image = np.array(predict_image, dtype=bool)
    # gt_image = np.array(gt_image, dtype=bool)

    # 计算True Positive（TP）
    tp = np.sum(np.logical_and(predict_image, gt_image))

    # 计算True Negative（TN）
    tn = np.sum(np.logical_and(np.logical_not(predict_image), np.logical_not(gt_image)))

    # 计算False Positive（FP）
    fp = np.sum(np.logical_and(predict_image, np.logical_not(gt_image)))

    # 计算False Negative（FN）
    fn = np.sum(np.logical_and(np.logical_not(predict_image), gt_image))

    # 计算IOU（Intersection over Union）
    iou = tp / (tp + fn + fp + 1e-7)

    # 计算Dice Coefficient（Dice系数）
    dice_coefficient = 2 * tp / (2 * tp + fn + fp + 1e-7)

    # 计算Accuracy（准确率）
    accuracy = (tp + tn) / (tp + fp + tn + fn + 1e-7)

    # 计算precision（精确率）
    precision = tp / (tp + fp + 1e-7)

    # 计算recall（召回率）
    recall = tp / (tp + fn + 1e-7)

    # 计算Sensitivity（敏感度）
    sensitivity = tp / (tp + fn + 1e-7)

    # 计算F1-score
    f1 = 2 * (precision * recall) / (precision + recall + 1e-7)

    # 计算Specificity（特异度）
    specificity = tn / (tn + fp + 1e-7)

    if evaluate == "iou":
        return iou

    if evaluate == "dice_coefficient":
        return dice_coefficient

    if evaluate == "accuracy":
        return accuracy

    if evaluate == "precision":
        return precision

    if evaluate == "recall":
        return recall

    if evaluate == "sensitivity":
        return sensitivity

    if evaluate == "f1":
        return f1

    if evaluate == "specificity":
        return specificity


def plot2d(gx, fx, fp, at=1, png=None):
    fig = plt.figure(figsize=(15, 5))
    # fig = plt.figure()
    ax = fig.add_subplot(131)
    ax.imshow(gx, vmin=-2, vmax=2, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    plt.savefig(pngDir + png + 'dd.png')
    ax = fig.add_subplot(132)
    ax.imshow(fx, vmin=0, vmax=1, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    ax = fig.add_subplot(133)
    ax.imshow(fp, vmin=0, vmax=1.0, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    if png:
        plt.savefig(pngDir + png + '.png')
    # cbar = plt.colorbar()
    # cbar.set_label('Fault probability')
    plt.tight_layout()
    plt.show()


def separate_plot2d(gx, fx, fp, at=1, png=None):
    plt.figure()
    plt.imshow(gx, vmin=-2, vmax=2, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    plt.show()
    plt.imshow(fx, vmin=0, vmax=1, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    plt.show()
    plt.imshow(fp, vmin=0, vmax=1.0, cmap=plt.cm.bone, interpolation='bicubic', aspect=at)
    plt.show()


if __name__ == '__main__':
    main()
